package com.ziiwee.reduce.rest.jpa;

import com.ziiwee.reduce.commons.utils.StringUtil;
import com.ziiwee.reduce.rest.Resource;
import com.ziiwee.reduce.rest.condition.Condition;
import com.ziiwee.reduce.rest.persistent.DataAccessor;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.data.jpa.repository.support.JpaEntityInformation;
import org.springframework.data.jpa.repository.support.JpaEntityInformationSupport;
import org.springframework.data.util.ProxyUtils;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import javax.persistence.EntityManager;
import javax.persistence.NamedEntityGraph;
import javax.persistence.TypedQuery;
import javax.persistence.criteria.*;
import javax.persistence.metamodel.EntityType;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;


@Component
@Transactional
public class JpaDataAccessor implements DataAccessor {

    public final static String HINT_FETCH_GRAPH = "javax.persistence.fetchgraph";

    protected final EntityManager entityManager;

    public JpaDataAccessor(EntityManager entityManager) {
        this.entityManager = entityManager;
    }


    @Override
    public boolean support(Class<?> entityClass) {
        for (EntityType<?> entity : entityManager.getMetamodel().getEntities()) {
            if (entity.getJavaType().equals(entityClass)) {
                return true;
            }
        }
        return false;
    }

    @Override
    public <T> long count(Class<T> entityClass) {
        return count(entityClass, null);
    }

    @Override
    public <T> long count(Class<T> entityClass, Condition condition) {
        CriteriaBuilder cb = entityManager.getCriteriaBuilder();
        CriteriaQuery<Long> query = cb.createQuery(Long.class);
        Root<T> root = query.from(entityClass);
        query.select(cb.count(root));
        Predicate predicate = toPredicate(condition, root, query, cb);
        if (predicate != null) {
            query.where(predicate);
        }
        TypedQuery<Long> typed = entityManager.createQuery(query);
        return typed.getSingleResult();
    }

    @Override
    public <T> List<T> findAll(Class<T> entityClass) {
        return findAll(entityClass, null, null);
    }

    @Override
    public <T> List<T> findAll(Class<T> entityClass, Sort sort) {
        return findAll(entityClass, null, sort);
    }

    @Override
    public <T> List<T> findAll(Class<T> entityClass, Condition condition) {
        return findAll(entityClass, condition, null);
    }

    @Override
    public <T> List<T> findAll(Class<T> entityClass, Condition condition, Sort sort) {
        TypedQuery<T> typed = customQuery(entityClass, (root, query, cb) -> {
            Predicate predicate = toPredicate(condition, root, query, cb);
            if (predicate != null) {
                query.where(predicate);
            }
            List<Order> orders = toOrder(sort, root, cb);
            if (!orders.isEmpty()) {
                query.orderBy(orders);
            }
        });
        return typed.getResultList();
    }

    @Override
    public <T> Page<T> findPage(Class<T> entityClass, Pageable pageable) {
        return findPage(entityClass, pageable, null);
    }

    @Override
    public <T> Page<T> findPage(Class<T> entityClass, Pageable pageable, Condition condition) {
        TypedQuery<T> typed = customQuery(entityClass, (root, query, cb) -> {
            Predicate predicate = toPredicate(condition, root, query, cb);
            if (predicate != null) {
                query.where(predicate);
            }
            if (pageable != null) {
                List<Order> orders = toOrder(pageable.getSort(), root, cb);
                if (!orders.isEmpty()) {
                    query.orderBy(orders);
                }
            }
        });
        if (pageable != null) {
            typed.setFirstResult(pageable.getPageNumber() * pageable.getPageSize());
            typed.setMaxResults(pageable.getPageSize());
            return new PageImpl<>(typed.getResultList(), pageable, count(entityClass, condition));
        }
        return new PageImpl<>(typed.getResultList());
    }

    @Override
    public <T> Optional<T> findOne(Class<T> entityClass, Object id) {
        return Optional.of(entityManager.find(entityClass, id));
    }

    @Override
    public <T> Optional<T> findOneByCondition(Class<T> entityClass, Condition condition) {
        TypedQuery<T> typed = customQuery(entityClass, (root, query, cb) -> {
            Predicate predicate = toPredicate(condition, root, query, cb);
            if (predicate != null) {
                query.where(predicate);
            }
        });
        List<T> resultList = typed.getResultList();
        if (resultList.size() > 0) {
            return Optional.of(resultList.get(0));
        }
        return Optional.empty();
    }

    @Override
    public <T> T save(T t) {
        if (getEntityInformation(t.getClass()).isNew(t)) {
            entityManager.persist(t);
            return t;
        } else {
            return entityManager.merge(t);
        }
    }

    @Override
    public <T> void delete(T t) {
        JpaEntityInformation<Object, ?> entityInformation = getEntityInformation(t.getClass());
        if (entityInformation.isNew(t)) {
            return;
        }
        Class<?> type = ProxyUtils.getUserClass(t);
        if (entityManager.find(type, entityInformation.getId(t)) == null) {
            return;
        }
        entityManager.remove(entityManager.contains(t) ? t : entityManager.merge(t));
    }

    @Override
    public <T> void deleteById(Class<T> entityClass, Object id) {
        findOne(entityClass, id).ifPresent(this::delete);
    }

    public interface CustomQuery<T> {
        void apply(Root<T> root, CriteriaQuery<T> query, CriteriaBuilder cb);
    }

    public <T> TypedQuery<T> customQuery(Class<T> entityClass, CustomQuery<T> customQuery) {
        CriteriaBuilder cb = entityManager.getCriteriaBuilder();
        CriteriaQuery<T> query = cb.createQuery(entityClass);
        Root<T> root = query.from(entityClass);
        customQuery.apply(root, query, cb);
        TypedQuery<T> Typed = entityManager.createQuery(query);
        String graphName = getGraphName(entityClass);
        if (StringUtil.hasLength(graphName)) {
            Typed.setHint(HINT_FETCH_GRAPH, entityManager.getEntityGraph(graphName));
        }
        return Typed;
    }

    public String getGraphName(Class<?> entityClass) {
        for (NamedEntityGraph graph : entityClass.getAnnotationsByType(NamedEntityGraph.class)) {
            if (StringUtil.hasLength(graph.name())) {
                return graph.name();
            }
        }
        return null;
    }

    public <T> Predicate toPredicate(Condition condition, Root<T> root, CriteriaQuery<?> query, CriteriaBuilder cb) {
        if (condition != null) {
            JpaConditionVisitor<T> visitor = new JpaConditionVisitor<>();
            Specification<T> spec = condition.accept(visitor);
            return spec == null ? null : spec.toPredicate(root, query, cb);
        }
        return null;
    }

    public <T> List<Order> toOrder(Sort sort, Root<T> root, CriteriaBuilder cb) {
        List<Order> orders = new ArrayList<>();
        if (sort == null) {
            return orders;
        }
        sort.get().forEach(it -> {
            Path<Object> path = root.get(it.getProperty());
            orders.add(it.getDirection().isAscending() ? cb.asc(path) : cb.desc(path));
        });
        return orders;
    }

    public <T> JpaEntityInformation<T, ?> getEntityInformation(Class<?> entityClass) {
        @SuppressWarnings("unchecked")
        Class<T> domainClass = (Class<T>) entityClass;
        return JpaEntityInformationSupport.getEntityInformation(domainClass, entityManager);
    }
}
